#ifndef SEGMENT_ANYTHING_MODEL_H
#define SEGMENT_ANYTHING_MODEL_H

#include <iostream>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <map>
#include <vector>
#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>

class SegmentAnythingModel
{
public:
	SegmentAnythingModel(const std::string& encoder_path, const std::string& decoder_path);
	~SegmentAnythingModel();

	void set_image(const cv::Mat& image_input);
	cv::Mat predict_mask_from_points(const std::vector<cv::Point2f>& points, const std::vector<int>& point_labels);
	cv::Mat predict_polygon_from_points(const std::vector<cv::Point2f>& points, const std::vector<int>& point_labels);

private:
	void compute_and_cache_image_embedding();
	cv::Mat get_image_embedding();

	int image_size_;

	Ort::Env env_;
	Ort::SessionOptions session_options_;
	Ort::Session encoder_session_;
	Ort::Session decoder_session_;

	std::mutex mtx_;
	std::condition_variable cv_;

	std::map<std::string, cv::Mat> image_embedding_cache_;

	std::thread worker_thread_;
	bool embedding_computed_;

	cv::Mat image_;
	cv::Mat image_embedding_;
};

#endif // SEGMENT_ANYTHING_MODEL_H
